#pragma once
#include <cmath>
#include <vector>
#include <algorithm>
#include <Common/headers.h>
#include <Utility/helper.hpp>
#include <Geometry/primitive.h>
#include <Geometry/model.h>
#include <Graphics/VirtualShader/fragment.hpp>

// jemalloc replacement
void* operator new(size_t size) {
	return malloc(size);
}

void operator delete (void* ptr) noexcept {
	free(ptr);
}

namespace occ {

	struct PolygonInfo {
		Plane plane;  // plane ax+by+cz+d=0
		int prim_id;               // polygon id
		int edge_id;
	};

	struct EdgeInfo {
		occFloat upper_x;     // x of the edge's upper endpoint
		occFloat dx;          // dx of intersections between neibouring scanlines
		int ymax;
		int dy;               // number of scanlines that the edge crossed
		int id;               // polygon id
	};

	struct ActiveEdgeInfo {
		occFloat xl, xr;
		occFloat dxl, dxr;
		int dyl, dyr;
		occFloat zl, dzx, dzy;
		int id;   // polygon id
		int next_edge;
	};

	struct ScanlineZBufferDescription {
		int w = 1024;
		int h = 768;
		vec3f light_pos = vec3f(0, 0, 0);
		CameraDescription* camera_desc = nullptr;
	};

	/* Original version withour parallelization */
	class ScanlineZBuffer {
	public:
		ScanlineZBuffer(ScanlineZBufferDescription* scanline_desc) :
			bg(DefaultColor::Black()),
			framebuffer(nullptr),
			pMesh(nullptr),
			pLocal2world(nullptr),
			pWorld2view(nullptr),
			light_color(DefaultColor::White())
		{
			h = scanline_desc->h;
			w = scanline_desc->w;
			CPT.resize(h);
			zbuffer.resize(w * h);

			perspective = getPerspectiveMatrix(
				scanline_desc->camera_desc->zNear,
				scanline_desc->camera_desc->zFar,
				scanline_desc->camera_desc->fovy,
				w, h);

			inverseZ = glm::identity<mat4f>();
			inverseZ[2].z = -1;

			light_pos = scanline_desc->light_pos;  // world coordinates
			//light_pos.z = -light_pos.z;

		}

		~ScanlineZBuffer() {
			//delete[] zlock_line;
		}
		
		void setLightPosition(const vec3f& pos) {
			light_pos = vec3f(pos.x, pos.y, pos.z);
		}

		void setMesh(TriMesh* pMesh) { this->pMesh = pMesh; }
		void setLocal2world(mat4f* pLocal2world) { this->pLocal2world = pLocal2world; }
		void setWorld2view(mat4f* pWorld2view) { this->pWorld2view = pWorld2view; }
		void setFramebuffer(Color* framebuffer) { this->framebuffer = framebuffer; }
		void setBackgroundColor(const Color& bg) { this->bg = bg; }
		void unbindMesh() { pMesh = nullptr; }
		void unbindLocal2world() { pLocal2world = nullptr; }
		void unbindWorld2view() { pWorld2view = nullptr; }
		void unbindFramebuffer() { framebuffer = nullptr; }

		// called by the user. the pipeline does not reset framebuffer automatically.
		void resetFramebuffer() {
			std::fill(framebuffer, framebuffer + h * w, bg); 
		}

		// called by the user. the pipeline does not reset zbuffer automatically.
		void resetZbuffer() {
			std::fill(zbuffer.begin(), zbuffer.end(), 0);  // set zbuffer to zFar side
		}

		void runPipeline() {
			assert(pMesh);
			assert(framebuffer);
			processVertex();
			processPixel();
			postprocess();
		}

		void drawPoints() {
			std::fill(framebuffer, framebuffer + h * w, bg);
			mat4f trans = perspective;
			if (pWorld2view) trans *= *pWorld2view;
			if (pLocal2world) trans *= *pLocal2world;
			const occFloat wh = w * 0.5;
			const occFloat hh = h * 0.5;
			for (int i = 0; i < pMesh->nVertices(); ++i) {
				vec4f p(pMesh->vp[i], 1);
				
				p = toNDC(trans, p);
				if (inNDC(p)) {
					p.x = wh * p.x + wh;
					p.y = hh * p.y + hh;
					int x = round(p.x);
					int y = round(p.y);
					if (0 <= x && x < w && 0 <= y && y < h) {
						framebuffer[w * y + x] = DefaultColor::Yellow();
					}

				}
			}
		}
	protected:
		void resetStatus() {
			/* clear previous data */
			tri_pvs.clear();
			pvs_idx.clear();
			fi.clear();
			CET.clear();
			for (int i = 0; i < h; ++i) {
				CPT[i].clear();
			}

			AET.clear();
			AET_idx.clear();
			AET_flag.clear();
		}
		void processVertex() {
			mat4f trans = perspective;
			if (pWorld2view) trans *= *pWorld2view;
			trans *= inverseZ;
			if (pLocal2world) trans *= *pLocal2world;

			/* pass 1: pt => NDC space => clipping => screen space */
			const occFloat wh = w * 0.5;
			const occFloat hh = h * 0.5;
			for (int i = 0; i < pMesh->nFaces(); ++i) {
				auto& face = pMesh->face[i];
				vec4f p1(pMesh->vp[face.x], 1);
				vec4f p2(pMesh->vp[face.y], 1);
				vec4f p3(pMesh->vp[face.z], 1);

				p1 = toNDC(trans, p1);
				p2 = toNDC(trans, p2);
				p3 = toNDC(trans, p3);
				if (inNDC(p1) || inNDC(p2) || inNDC(p3)) {
					// transform to screen space
					p1.x = wh * p1.x + wh;
					p2.x = wh * p2.x + wh;
					p3.x = wh * p3.x + wh;
					p1.y = wh * p1.y + wh;
					p2.y = wh * p2.y + wh;
					p3.y = wh * p3.y + wh;

					// use 1-z for better precision. see RealTimeRendering 4th, chap. 4
					p1.z = 1-(p1.z + 1.0) * 0.5;
					p2.z = 1-(p2.z + 1.0) * 0.5;
					p3.z = 1-(p3.z + 1.0) * 0.5;


					tri_pvs.emplace_back(p1, p2, p3);
					fi.emplace_back(i);
				}
			}

			//pvs_idx.resize(tri_pvs.size());
			//std::iota(pvs_idx.begin(), pvs_idx.end(), 0);
			//std::sort(pvs_idx.begin(), pvs_idx.end(), InverseZCompareIdx(tri_pvs.data()));
			CET.reserve(tri_pvs.size() * 3);

			/* pass 2: construct classified polygon and edge tables */ 
			for (int i = 0; i < tri_pvs.size(); ++i) {
				//int idx = pvs_idx[i];
				int idx = i;
				int original_tri_idx = fi[idx];
				auto& tri = tri_pvs[idx];
				/* construct classified polygon table */
				PolygonInfo polyinfo = { Plane(tri.v1,tri.v2,tri.v3),original_tri_idx, i };
				int ymax = round(max(tri.v1.y, tri.v2.y, tri.v3.y));
				int ymin = round(min(tri.v1.y, tri.v2.y, tri.v3.y));
				bool up = ymax >= h;  // part of the triangle is above the screen
				if (ymax >= 0 && ymax - ymin > 0 && abs(polyinfo.plane.c) >= EPS) {
					if (up)
						CPT[h - 1].emplace_back(polyinfo);
					else CPT[ymax].emplace_back(polyinfo);
				}

				/* construct classified edge tables */
				EdgeInfo e[3];
				int ei0 = 0, ei1 = 1, ei2 = 2;
				for (int j = 0; j < 3; ++j) {
					vec3f* pu = &tri._v[j], * pl = &tri._v[(j + 1) % 3];
					if (pu->y < pl->y) pu = &tri._v[(j + 1) % 3], pl = &tri._v[j];
					occFloat dx = pu->x - pl->x;
					int eymax = round(pu->y);
					int eymin = round(pl->y);
					int dy = eymax - eymin;
					dx = dy != 0 ? -dx / (pu->y - pl->y) : 0;
					if (eymax >= h) {
						dy = h - 1 - eymin;
						occFloat alpha = (pu->y - h) / (pu->y - pl->y);
						e[j].upper_x = pu->x * (1.0 - alpha) + pl->x * alpha;
					}
					else e[j].upper_x = pu->x;

					if (eymax != ymax || dy == 0) {
						ei2 = j;
						ei0 = (j + 1) % 3;
						ei1 = (j + 2) % 3;
					}

					e[j].dx = dx;
					e[j].dy = dy;
					e[j].id = original_tri_idx;
					e[j].ymax = eymax;
				}
				if (e[ei0].upper_x > e[ei1].upper_x) {
					int t = ei0;
					ei0 = ei1;
					ei1 = t;
				}
				// ensure sequence: 
				//top left edge, top right edge, bottom edge/ top horizontal edge
				CET.emplace_back(e[ei0]);
				CET.emplace_back(e[ei1]);
				CET.emplace_back(e[ei2]);
			}
		}

		void processPixel() {
			//mat4f trans = inverseZ;
			mat4f trans = glm::identity<mat4f>();
			if (pLocal2world) trans *= *pLocal2world;
			//if (pWorld2view) trans *= *pWorld2view;
			//mat4f trans_normal = glm::transpose(glm::inverse(trans));
			mat4f trans_normal = trans;

			for (int i = h-1; i >= 0; --i) {
				/* add active pair sets when finding new classified polygons */
				for (auto& polyinfo : CPT[i]) {
					size_t id = 3 * polyinfo.edge_id;
					auto& e0 = CET[id + 0];
					auto& e1 = CET[id + 1];
					ActiveEdgeInfo ae;
					ae.xl = e0.upper_x;
					ae.dxl = e0.dx;
					ae.dyl = e0.dy;
					ae.xr = e1.upper_x;
					ae.dxr = e1.dx;
					ae.dyr = e1.dy;
					occFloat a = polyinfo.plane.a;
					occFloat b = polyinfo.plane.b;
					occFloat c = polyinfo.plane.c;
					occFloat d = polyinfo.plane.d;
					ae.dzx = -a / c;
					ae.dzy = b / c;
					ae.zl = (-d - a * ae.xl - b * i) / c;
					ae.id = polyinfo.prim_id;
					ae.next_edge = CET[id + 2].dy != 0 ? id + 2 : -1;
					AET.emplace_back(ae);
					AET_flag.push_back(1);
				}
				/* update AET idx */
				AET_idx.clear();
				for (int k = 0; k < AET_flag.size(); ++k) {
					if (AET_flag[k]) AET_idx.push_back(k);
				}
				// possible for parallelization
				for (int k = 0; k < AET_idx.size(); ++k) {
					int idx = AET_idx[k];
					auto& edgepair = AET[idx];
					occFloat zx = edgepair.zl;
					int xl = round(edgepair.xl);
					int xr = round(edgepair.xr);

					/* incremental depth update */
					if (xl < 0) {
						zx += (-xl) * edgepair.dzx;
						xl = 0;
					}
					for (int j = xl; j <= xr; ++j) {
						if (j >= w) {
							break;
						}
						int pix = w * i + w - j;
						if (zx > zbuffer[pix]) {
							zbuffer[pix] = zx;
							vec4f center_world = trans * vec4f(pMesh->face_center[edgepair.id], 1);
							vec4f normal_world = glm::normalize(trans_normal * vec4f(pMesh->face_normal[edgepair.id], 0));
							vec4f _light_pos = vec4f(light_pos, 1);
							//if (pWorld2view) _light_pos = *pWorld2view * _light_pos;
							Color color;
							faceFragmentShader(center_world, normal_world, _light_pos, light_color, pMesh->color, color);
							//facePhongFragmentShader(center_world, normal_world, _light_pos, light_color, pMesh->color, color);
							framebuffer[pix] = color;
						}
	
						zx += edgepair.dzx;
					}
					/* update active edge pair */
					edgepair.dyl--;
					edgepair.dyr--;
					if ((edgepair.dyl) <= 0 && (edgepair.dyr) <= 0) {
						AET_flag[idx] = 0;  // lazy delete
						continue;
					}
					edgepair.xl += edgepair.dxl;

					edgepair.xr += edgepair.dxr;
					if ((edgepair.dyl <= 0 || edgepair.dyr <= 0) && edgepair.next_edge >= 0) {
						auto& e = CET[edgepair.next_edge];
						edgepair.next_edge = -1;
						if (edgepair.dyl == 0) {
							edgepair.xl = e.upper_x;
							edgepair.dxl = e.dx;
							edgepair.dyl = e.dy;
						}
						if (edgepair.dyr == 0) {
							edgepair.xr = e.upper_x;
							edgepair.dxr = e.dx;
							edgepair.dyr = e.dy;
						}
					}
					edgepair.zl += edgepair.dxl * edgepair.dzx + edgepair.dzy;
				}
				
			}
		}

		void postprocess() { resetStatus(); }

	protected:
		int h, w;
		TriMesh* pMesh;
		mat4f* pLocal2world;
		mat4f* pWorld2view;
		Color* framebuffer;
		Color bg;
		Color light_color;
		vec3f light_pos;

		mat4f perspective;
		mat4f inverseZ;
		std::vector<Triangle> tri_pvs;  // potential visible set in sreen space
		std::vector<int> pvs_idx;
		std::vector<int> fi;  // face idx of pvs
		std::vector<std::vector<PolygonInfo>> CPT;  // classified polygon table
		std::vector<EdgeInfo> CET;                  // classified edge table
		std::vector<ActiveEdgeInfo> AET;            // active edge table
		std::vector<int> AET_flag;
		std::vector<int> AET_idx;
		std::vector<occFloat> zbuffer;
	};
}